import sys
import logging
import copy
import torch
from learners import get_learner
from utils.data.data_manager import DataManager
from utils.toolkit import count_parameters
import os
from datetime import datetime
import shutil
import torch.backends
import torch.backends.cudnn


def train(args):
    seed_list = copy.deepcopy(args["seed"])
    device = copy.deepcopy(args["device"])

    for seed in seed_list:
        args["seed"] = seed
        args["device"] = device
        if args.get("train_seed", None) is None:
            print("========== No train_seed in config, use seed instead. ==========")
            args["train_seed"] = args["seed"]
        _train(args)


def _train(args):
    # FIXME: This is a temporary solution to load the ViT models for fine-tuning
    assert isinstance(args["seed"], int), "seed must be an integer"
    # init_cls = 0 if args["init_cls"] == args["inc_cls"] else args["init_cls"]
    # assert init_cls == 0, "init_cls must be 0 for incremental learning"

    current_time = datetime.now().strftime("%Y%m%d%H%M%S")
    current_time = args["prefix"] if args["prefix"] else current_time
    logfile_path = "logs/{}/{}/{}".format(
        args["learner_name"],
        args["dataset"],
        current_time,
    )
    if not os.path.exists(logfile_path):
        os.makedirs(logfile_path)
    logfilename = os.path.join(
        logfile_path,
        "{}_{}_{}".format(
            args["prefix"],
            args["seed"],
            args["backbone_type"],
        ),
    )

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(filename)s] => %(message)s",
        handlers=[
            logging.FileHandler(filename=logfilename + ".log"),
            logging.StreamHandler(sys.stdout),
        ],
    )

    suffix = args["backbone_type"].split("_")[-1]
    shutil.copy2(args["config"], logfile_path)
    shutil.copy2(f"learners/{args['learner_name']}.py", logfile_path)
    try:
        shutil.copy2(f"backbone/vit_{suffix}.py", logfile_path)
    except FileNotFoundError:
        print(f"backbone/vit_{suffix}.py not found, using backbone/vit.py instead.")
        shutil.copy2("backbone/vit.py", logfile_path)

    _set_random(args["train_seed"])
    _set_device(args)
    print_args(args)

    ckp_path = "_ckps/{}/{}/{}/{}".format(
        args["learner_name"],
        args["backbone_type"],
        args["dataset"],
        current_time,
    )
    args["ckp_path"] = ckp_path if getattr(args, "save_ckp", False) else None

    data_manager = DataManager(
        args["dataset"],
        args["shuffle"],
        args["seed"],
        args["init_cls"],
        args["inc_cls"],
        args,
    )

    args["nb_classes"] = data_manager.nb_classes  # update args
    args["nb_tasks"] = data_manager.nb_tasks
    model = get_learner(args["learner_name"], args, data_manager)
    learner_topk = f"top{model.topk}"
    cnn_curve, nme_curve = {"top1": [], learner_topk: []}, {
        "top1": [],
        learner_topk: [],
    }
    for task in range(data_manager.nb_tasks):
        logging.info("All params: {}".format(count_parameters(model._network)))
        logging.info(
            "Trainable params: {}".format(count_parameters(model._network, True))
        )

        model.incremental_train()
        cnn_accy, nme_accy = model.eval_task()

        if nme_accy is not None:
            logging.info("CNN: {}".format(cnn_accy["grouped"]))
            logging.info("NME: {}".format(nme_accy["grouped"]))

            cnn_curve["top1"].append(cnn_accy["top1"])
            cnn_curve[learner_topk].append(cnn_accy[learner_topk])

            nme_curve["top1"].append(nme_accy["top1"])
            nme_curve[learner_topk].append(nme_accy[learner_topk])

            logging.info("CNN top1 curve: {}".format(cnn_curve["top1"]))
            logging.info(
                "CNN {} curve: {}".format(learner_topk, cnn_curve[learner_topk])
            )
            logging.info("NME top1 curve: {}".format(nme_curve["top1"]))
            logging.info(
                "NME {} curve: {}\n".format(learner_topk, nme_curve[learner_topk])
            )

            print(
                "Average Accuracy (CNN):",
                sum(cnn_curve["top1"]) / len(cnn_curve["top1"]),
            )
            print(
                "Average Accuracy (NME):",
                sum(nme_curve["top1"]) / len(nme_curve["top1"]),
            )

            logging.info(
                "Average Accuracy (CNN): {}".format(
                    sum(cnn_curve["top1"]) / len(cnn_curve["top1"])
                )
            )
            logging.info(
                "Average Accuracy (NME): {}".format(
                    sum(nme_curve["top1"]) / len(nme_curve["top1"])
                )
            )
        else:
            logging.info("No NME accuracy.")
            logging.info("CNN: {}".format(cnn_accy["grouped"]))

            cnn_curve["top1"].append(cnn_accy["top1"])
            cnn_curve[learner_topk].append(cnn_accy[learner_topk])

            logging.info("CNN top1 curve: {}".format(cnn_curve["top1"]))
            logging.info(
                "CNN {} curve: {}\n".format(learner_topk, cnn_curve[learner_topk])
            )

            print(
                "Average Accuracy (CNN):",
                sum(cnn_curve["top1"]) / len(cnn_curve["top1"]),
            )
            logging.info(
                "Average Accuracy (CNN): {} \n".format(
                    sum(cnn_curve["top1"]) / len(cnn_curve["top1"])
                )
            )


def _set_device(args):
    device_type = args["device"]
    gpus = []

    for device in device_type:
        if device == -1:
            device = torch.device("cpu")
        else:
            device = torch.device("cuda:{}".format(device))

        gpus.append(device)

    args["device"] = gpus


def _set_random(seed=1):
    print("Setting training seed to {}".format(seed))
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def print_args(args):
    for key, value in args.items():
        logging.info("{}: {}".format(key, value))
